# 加载 circlize 等必备的包
library(circlize)
library(gridBase)
library(ComplexHeatmap)

# 生成模拟数据集
type = c(rep("Tumor", 10), rep("Control", 10))

set.seed(888)

######################################
# generate methylation matrix
rand_meth = function(k, mean) {
  (runif(k) - 0.5)*min(c(1-mean), mean) + mean
}

mean_meth = c(rand_meth(300, 0.3), rand_meth(700, 0.7))
mat_meth = as.data.frame(lapply(mean_meth, function(m) {
  if(m < 0.3) {
    c(rand_meth(10, m), rand_meth(10, m + 0.2))
  } else if(m > 0.7) {
    c(rand_meth(10, m), rand_meth(10, m - 0.2))
  } else {
    c(rand_meth(10, m), rand_meth(10, m + sample(c(1, -1), 1)*0.2))
  }

}))
mat_meth = t(mat_meth)
rownames(mat_meth) = NULL
colnames(mat_meth) = paste0("sample", 1:20)

######################################
# generate directions for methylation
direction = rowMeans(mat_meth[, 1:10]) - rowMeans(mat_meth[, 11:20])
direction = ifelse(direction > 0, "hyper", "hypo")

#######################################
# generate expression matrix
mat_expr = t(apply(mat_meth, 1, function(x) {
  x = x + rnorm(length(x), sd = (runif(1)-0.5)*0.4 + 0.5)
  -scale(x)
}))
dimnames(mat_expr) = dimnames(mat_meth)

#############################################################
# matrix for correlation between methylation and expression
cor_pvalue = -log10(sapply(seq_len(nrow(mat_meth)), function(i) {
  cor.test(mat_meth[i, ], mat_expr[i, ])$p.value
}))

#####################################################
# matrix for types of genes
gene_type = sample(c("protein_coding", "lincRNA", "microRNA", "psedo-gene", "others"),
                   nrow(mat_meth), replace = TRUE, prob = c(6, 1, 1, 1, 1))

#################################################
# annotation to genes
anno_gene = sapply(mean_meth, function(m) {
  if(m > 0.6) {
    if(runif(1) < 0.8) return("intragenic")
  }
  if(m < 0.4) {
    if(runif(1) < 0.4) return("TSS")
  }
  return("intergenic")
})

############################################
# distance to genes
dist = sapply(mean_meth, function(m) {
  if(m < 0.6) {
    if(runif(1) < 0.8) return(round( (runif(1)-0.5)*1000000 + 500000 ))
  }
  if(m < 0.3) {
    if(runif(1) < 0.4) return(round( (runif(1) - 0.5)*1000 + 500))
  }
  return(round( (runif(1) - 0.5)*100000 + 50000))
})


#######################################
# annotation to enhancers
rand_enhancer = function(m) {
  if(m < 0.4) {
    if(runif(1) < 0.6) return(runif(1))
  } else if (runif(1) < 0.1) {
    return(runif(1))
  }
  return(0)
}
anno_enhancer_1 = sapply(mean_meth, rand_enhancer)
anno_enhancer_2 = sapply(mean_meth, rand_enhancer)
anno_enhancer_3 = sapply(mean_meth, rand_enhancer)
anno_enhancer = data.frame(enhancer_1 = anno_enhancer_1, enhancer_2 = anno_enhancer_2, enhancer_3 = anno_enhancer_3)

#################################
# put everything into one object
res_list = list()
res_list$type = type
res_list$mat_meth = mat_meth
res_list$mat_expr = mat_expr
res_list$direction = direction
res_list$cor_pvalue = cor_pvalue
res_list$gene_type = gene_type
res_list$anno_gene = anno_gene
res_list$dist = dist
res_list$anno_enhancer = anno_enhancer

# 设置cluster
set.seed(123)
km = kmeans(mat_meth, centers = 5)$cluster

# 绘制基础班环形热图
col_meth = colorRamp2(c(0, 0.5, 1), c("blue", "white", "red"))
circos.heatmap(mat_meth, split = km, col = col_meth, track.height = 0.12)

col_direction = c("hyper" = "red", "hypo" = "blue")
circos.heatmap(direction, col = col_direction, track.height = 0.01)

col_expr = colorRamp2(c(-2, 0, 2), c("green", "white", "red"))
circos.heatmap(mat_expr, col = col_expr, track.height = 0.12)

col_pvalue = colorRamp2(c(0, 2, 4), c("white", "white", "red"))
circos.heatmap(cor_pvalue, col = col_pvalue, track.height = 0.01)

library(RColorBrewer)
col_gene_type = structure(brewer.pal(length(unique(gene_type)), "Set3"), names = unique(gene_type))
circos.heatmap(gene_type, col = col_gene_type, track.height = 0.01)

col_anno_gene = structure(brewer.pal(length(unique(anno_gene)), "Set1"), names = unique(anno_gene))
circos.heatmap(anno_gene, col = col_anno_gene, track.height = 0.01)

col_dist = colorRamp2(c(0, 10000), c("black", "white"))
circos.heatmap(dist, col = col_dist, track.height = 0.01)

col_enhancer = colorRamp2(c(0, 1), c("white", "orange"))
circos.heatmap(anno_enhancer, col = col_enhancer, track.height = 0.03)




## 绘制带有链接线的环形热图

## 在上图的基础上添加环线

df_link = data.frame(
  from_index = sample(nrow(mat_meth), 20),
  to_index = sample(nrow(mat_meth), 20)
)

for(i in seq_len(nrow(df_link))) {
  # Let's call the DMR with index df_link$from_index[i] as DMR1,
  # and the other one with index df_link$to_index[i] as DMR2.

  # The sector where DMR1 is in.
  group1 = km[ df_link$from_index[i] ]
  # The sector where DMR2 is in.
  group2 = km[ df_link$to_index[i] ]

  # The subset of DMRs (row indices from mat_meth) in sector `group1`.
  subset1 = get.cell.meta.data("subset", sector.index = group1)
  # The row ordering in sector `group1`.
  row_order1 = get.cell.meta.data("row_order", sector.index = group1)
  # This is the position of DMR1 in the `group1` heatmap.
  x1 = which(subset1[row_order1] == df_link$from_index[i])

  # The subset of DMRs (row indices from mat_meth) in sector `group2`.
  subset2 = get.cell.meta.data("subset", sector.index = group2)
  # The row ordering in sector `group2`.
  row_order2 = get.cell.meta.data("row_order", sector.index = group2)
  # This is the position of DMR2 in the `group2` heatmap.
  x2 = which(subset2[row_order2] == df_link$to_index[i])

  # We take the middle point and draw a link between DMR1 and DMR2
  circos.link(group1, x1 - 0.5, group2, x2 - 0.5, col = rand_color(1))
}

for(i in seq_len(nrow(df_link))) {
  circos.heatmap.link(df_link$from_index[i],
                      df_link$to_index[i],
                      col = rand_color(1))
}


## 绘制带有图例的环形热图

circos.clear()

circlize_plot = function() {
  circos.heatmap(mat_meth, split = km, col = col_meth, track.height = 0.12)
  circos.heatmap(direction, col = col_direction, track.height = 0.01)
  circos.heatmap(mat_expr, col = col_expr, track.height = 0.12)
  circos.heatmap(cor_pvalue, col = col_pvalue, track.height = 0.01)
  circos.heatmap(gene_type, col = col_gene_type, track.height = 0.01)
  circos.heatmap(anno_gene, col = col_anno_gene, track.height = 0.01)
  circos.heatmap(dist, col = col_dist, track.height = 0.01)
  circos.heatmap(anno_enhancer, col = col_enhancer, track.height = 0.03)

  for(i in seq_len(nrow(df_link))) {
    circos.heatmap.link(df_link$from_index[i],
                        df_link$to_index[i],
                        col = rand_color(1))
  }
  circos.clear()
}

lgd_meth = Legend(title = "Methylation", col_fun = col_meth)
lgd_direction = Legend(title = "Direction", at = names(col_direction),
                       legend_gp = gpar(fill = col_direction))
lgd_expr = Legend(title = "Expression", col_fun = col_expr)
lgd_pvalue = Legend(title = "P-value", col_fun = col_pvalue, at = c(0, 2, 4),
                    labels = c(1, 0.01, 0.0001))
lgd_gene_type = Legend(title = "Gene type", at = names(col_gene_type),
                       legend_gp = gpar(fill = col_gene_type))
lgd_anno_gene = Legend(title = "Gene anno", at = names(col_anno_gene),
                       legend_gp = gpar(fill = col_anno_gene))
lgd_dist = Legend(title = "Dist to TSS", col_fun = col_dist,
                  at = c(0, 5000, 10000), labels = c("0kb", "5kb", "10kb"))
lgd_enhancer = Legend(title = "Enhancer overlap", col_fun = col_enhancer,
                      at = c(0, 0.25, 0.5, 0.75, 1), labels = c("0%", "25%", "50%", "75%", "100%"))

plot.new()
circle_size = unit(1, "snpc") # snpc unit gives you a square region

pushViewport(viewport(x = 0, y = 0.5, width = circle_size, height = circle_size,
                      just = c("left", "center")))
par(omi = gridOMI(), new = TRUE)
circlize_plot()
upViewport()

h = dev.size()[2]
lgd_list = packLegend(lgd_meth, lgd_direction, lgd_expr, lgd_pvalue, lgd_gene_type,
                      lgd_anno_gene, lgd_dist, lgd_enhancer, max_height = unit(0.9*h, "inch"))
draw(lgd_list, x = circle_size, just = "left")



